# app/detection.py
import cv2
from ultralytics import solutions
from collections import defaultdict
from app.config import COCO_NAMES

class ObjectDetector:
    def __init__(self, region_points, model_path, conf=0.6, line_width=2):
        self.region_points = region_points
        self.model_path = model_path
        self.conf = conf
        self.line_width = line_width
        self.trackzone = None
        
        # 跟踪计数变量
        self.unique_object_ids = set()
        self.class_counts = defaultdict(int)
        self.id_to_class = {}
        self.frame_count = 0
        
    def initialize(self):
        self.trackzone = solutions.TrackZone(
            show=False,
            conf=self.conf,
            region=self.region_points,
            model=self.model_path,
            line_width=self.line_width,
        )
        self.reset_counters()
        
    def reset_counters(self):
        self.unique_object_ids = set()
        self.class_counts = defaultdict(int)
        self.id_to_class = {}
        self.frame_count = 0
        
    def process_frame(self, frame):
        self.frame_count += 1
        
        # 使用trackzone处理当前帧
        results = self.trackzone.process(frame)
        
        # 检查是否有ID和类别信息
        if hasattr(results, 'total_tracks') and results.total_tracks > 0:
            # 获取当前帧中的跟踪ID和类别
            for track_id, cls in zip(self.trackzone.track_ids, self.trackzone.clss):
                # 如果是新出现的ID，则计数加一
                if track_id not in self.unique_object_ids:
                    self.unique_object_ids.add(track_id)
                    class_id = int(cls)
                    class_name = COCO_NAMES.get(class_id, f"class_{class_id}")
                    self.class_counts[class_name] += 1
                    self.id_to_class[track_id] = class_name
        
        # 绘制统计信息
        stats_frame = results.plot_im.copy()
        y_offset = 40
        cv2.putText(stats_frame, f"Total Objects: {len(self.unique_object_ids)}", 
                    (20, y_offset), cv2.FONT_HERSHEY_SIMPLEX, 1, (0, 0, 255), 2)
        
        return stats_frame, len(self.unique_object_ids), self.class_counts